using LinearAlgebra
using Ranges
using Interpolations
using Random, Distributions
using Statistics

""" Grids """
function asset_grid(amin, amax, n)
    # find maximum ubar of uniform grid corresponding to desired maximum amax of asset grid
    ubar = log(1 + log(1 + amax - amin))

    # make uniform grid
    u_grid = range(0, ubar, n)

    # double-exponentiate uniform grid and add amin to get grid from amin to amax
    return amin .+ exp.(exp.(u_grid) .- 1) .- 1
end # works

""" Create grid between amin and amax equidistant in logs """
function agrid(amax, n; amin=0)
    pivot = abs(amin) + 0.25
    a_grid = exp10.(range(log10(amin + pivot), log10(amax + pivot), length=Int(n))) .- pivot
    a_grid[1] = amin # ensures equal to main
    return a_grid
end # works

# Temporarily include the old way of constructing grids from ikc_old for comparability of results
function agrid_old(amax, N, amin=0, frac=1/25)
    apivot = (1 + amin)^(1 - frac)*(1 + amax)^frac - 1
    a_old = exp10.(range(log10(amin + apivot), log10(amax + apivot), length=N)) .- apivot
    a_old[1] = amin  # make sure *exactly* equal to amin
    return a_old
end # works

""" Create grid between amin and amax. phi=1 equidistant, phi>1 dense near amin.
Extra flexibility useful in non-convex problems in which policy functions have
nonlinear and even non-monotonic sections far from the borrowing limit. """
function nonlinspace(amax, n, phi, amin=0)
    a_grid = zeros(n) #array of zeros
    a_grid[1] = amin #stores amin value to first entry of agrid
    for i in 2:n
        a_grid[i] = a_grid[i-1] + (amax - a_grid[i-1]) / (n-i)^phi
    end
    return a_grid
end # works

""" Markov Chains """

""" Find invariant dist of Markov chain by iteration """
function stationary(Pi::AbstractArray, pi_seed=nothing, tol=1E-11, maxit=10000)
    if isnothing(pi_seed)
        _pi = ones(size(Pi,1)) ./ size(Pi,1)
    else
        _pi = pi_seed
    end

    does_break = false
    for it in 1:maxit #range of max iterations
        pi_new = Pi' * _pi # matrix multiply
        if maximum(abs.(pi_new - _pi)) < tol
            does_break = true
            break
        end
        _pi = pi_new
    end

    if !does_break
        error("No convergence after max forward iterations") # gives error if no convergence
    end

    return _pi
end # works

""" Mean of discretized random variable with support x and probability mass function pi """
function mean(x, _pi) #mean
    return sum(_pi .* x) # works
end

""" Variance of discretized random variable with support x and probability mass function _pi """
function variance(x, _pi) #variance
    return sum(_pi .* (x .- sum(_pi .* x)).^2) # works
end

""" Standard deviation of discretized random variable with support x and probability mass function pi """
function std(x, _pi)#standard dev
    return sqrt(variance(x, _pi)) # works
end

""" Covariance of two discretized random variables with supports x and y common pmf pi """
function cov(x,y,_pi) #covariance
    return sum(_pi .* (x .- mean(x, _pi)) .* (y .- mean(y,_pi))) # works
end

""" Correlation of two discretized random variables with supports x and y common pmf pi """
function corr(x,y,_pi) #correlation
    return cov(x,y,_pi) / ((std(x,_pi) * std(y,_pi))) # works
end

function markov_tauchen(rho, sigma, N=7, m=3, normalize=true)
"""Tauchen method for discretizing AR(1) s_t = rho*s_(t-1) + eps_t.

    Parameters
    ----------
    rho   : scalar, persistence
    sigma : scalar, unconditional sd of s_t
    N     : int, number of states in discretized Markov process
    m     : scalar, discretized s goes from approx -m*sigma to m*sigma

    Returns
    ----------
    y  : array (N), states proportional to exp(s) s.t. E[y] = 1
    pi : array (N), stationary distribution of discretized process
    Pi : array (N*N), Markov matrix for discretized process
    """

    # make normalized grid starting with cross-sectional sd of 1
    s = range(-m, m, N) #look at how parameters are passed
    ds = s[2] - s[1]
    sd_innov = sqrt(1 - rho^2)

    # standard Tauchen method to generate Pi given N and m
    Pi =Array{Float64}(undef, N, N) # converts Pi to an array of dim NxN
    Pi[:,1] = cdf.(Normal(0, sd_innov), s[1] .- rho .* s .+ ds / 2)
    Pi[:,end] = 1 .- cdf.(Normal(0, sd_innov), s[end] .- rho .* s .- ds / 2)

    for j in 2:N-1
        Pi[:,j] = cdf.(Normal(0,sd_innov), s[j] .- rho .* s .+ ds / 2) .- cdf.(Normal(0,sd_innov),s[j] .- rho .* s .- ds / 2)
    end

    # invariant distribution and scaling
    _pi = stationary(Pi) # stationary distribution
    s *= (sigma / sqrt(variance(s, _pi)))

    if normalize
        y = exp.(s) ./ sum(_pi .* exp.(s)) # normalizing y
    else
        y = s
    end
    return y, _pi, Pi
end

function markov_rouwenhorst(rho, sigma, N=7)
    """Rouwenhorst method analog to markov_tauchen"""

    # Explicitly typecast N as an integer, since when the grid constructor functions
    # (e.g. the function that makes a_grid) are implemented as blocks, they interpret the
    # integer-valued calibration as a float.

    # parametrize Rouwenhorst for n=2
    N = Int(N) # convert N to integer
    p = (1 + rho) / 2
    Pi = [p 1-p; 1-p p]

    # implement recursion to build from n=3 to n=N
    for n in 3:N
        P1, P2, P3, P4 = (zeros(n,n) for _ in 1:4)
        P1[1:end-1, 1:end-1] = p .* Pi
        P2[1:end-1, 2:end] = (1 - p) .* Pi
        P3[2:end, 1:end-1] = (1 - p) .* Pi
        P4[2:end, 2:end] = p .* Pi
        Pi = P1 + P2 + P3 + P4
        Pi[2:end-1,:] = Pi[2:end-1,:] ./ 2
    end

    # invariant distribution and scaling
    _pi = stationary(Pi)
    s = range(-1, 1, N)
    s = s * (sigma / sqrt(variance(s, _pi)))
    y = exp.(s) ./ sum(_pi .* exp.(s))
    return y, _pi, Pi
end
